from typing import Optional, List

from peewee import Model, CharField, AutoField
from playhouse.shortcuts import model_to_dict
from pydantic import BaseModel

from db.my_sqlite import DB

# 定义用户模型
class User(Model):
    id = AutoField(primary_key=True)
    username = CharField(unique=True)
    password = CharField()
    email = CharField(null=True)
    role = CharField()
    api_key = CharField(null=True, unique=True)


    class Meta:
        database = DB

class UserModel(BaseModel):
    id: Optional[int]
    username: Optional[str]
    password: Optional[str]
    email: Optional[str]= ""
    role: Optional[str] = "user"
    api_key: Optional[str]


class UsersTable:
    # 创建表
    def __init__(self, db):
        self.db = db
        self.db.create_tables([User])

    @staticmethod
    def insert_new_user(
        username: str,
        password: str,
        email: str = "",
        role: str = "user",
        api_key: str = None,
    ) -> Optional[UserModel]:
        result = User.create(
            **{
                "username": username,
                "password": password,
                "email": email,
                "role": role,
                "api_key": api_key,
            }
        )
        if result:
            return Users.get_user_by_name(username)
        else:
            return None

    @staticmethod
    def get_user_by_id(id: int) -> Optional[UserModel]:
        try:
            user = User.get(User.id == id)
            return UserModel(**model_to_dict(user))
        except:
            return None

    @staticmethod
    def get_user_by_name(username: str) -> Optional[UserModel]:
        try:
            user = User.get(User.username == username)
            return UserModel(**model_to_dict(user))
        except:
            return None

    @staticmethod
    def get_user_by_api_key(api_key: str) -> Optional[UserModel]:
        try:
            user = User.get(User.api_key == api_key)
            return UserModel(**model_to_dict(user))
        except:
            return None

    @staticmethod
    def get_user_by_email(email: str) -> Optional[UserModel]:
        try:
            user = User.get(User.email == email)
            return UserModel(**model_to_dict(user))
        except:
            return None

    @staticmethod
    def get_users(skip: int = 0, limit: int = 50) -> List[UserModel]:
        return [
            UserModel(**model_to_dict(user))
            for user in User.select()
            .limit(limit).offset(skip)
        ]

    @staticmethod
    def get_num_users() -> Optional[int]:
        return User.select().count()

    @staticmethod
    def update_user_role_by_id(id: int, role: str) -> Optional[UserModel]:
        try:
            query = User.update(role=role).where(User.id == id)
            query.execute()

            user = User.get(User.id == id)
            return UserModel(**model_to_dict(user))
        except:
            return None

    @staticmethod
    def update_user_by_id(id: int, updated: dict) -> Optional[UserModel]:
        try:
            query = User.update(**updated).where(User.id == id)
            query.execute()

            user = User.get(User.id == id)
            return UserModel(**model_to_dict(user))
        except:
            return None

    @staticmethod
    def update_user_api_key_by_id(id: int, api_key: str) -> bool:
        try:
            query = User.update(api_key=api_key).where(User.id == id)
            result = query.execute()

            return True if result == 1 else False
        except:
            return False

    @staticmethod
    def get_user_api_key_by_id(id: int) -> Optional[str]:
        try:
            user = User.get(User.id == id)
            return user.api_key
        except:
            return None


Users = UsersTable(DB)

